#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Performs data capture for the Audio2Face-3D Authoring service hosted on NVCF."""

import argparse

import scipy

import auth
from common import (
    check_health,
    get_avatar_face_pose,
    make_face_pose_request,
    perform_parallel_data_exchange,
    positive_int,
    upload_audio_clip_and_get_hash,
    GrpcChannelParams,
    EMOTION_KEYS,
    TIME_1_FRAME,
)

from nvidia_ace.services.a2f_authoring.v1_pb2_grpc import A2FAuthoringServiceStub

PERF_DATA_FILE = "output_perf.txt"
OUT_IMG = "output_latency.png"
OUTPUT_BLENDSHAPE = "output_blendshape.csv"
OUTPUT_EMOTIONS = "output_emotions.csv"


def get_blendshapes_for_each_frame(audio_clip: str, metadata_args: list) -> tuple[list, list]:
    """
    Use the A2F-3D Authoring microservice to get blendshapes for each frame.

    Args:
      audio_clip (str): The path of the audio clip.
      metadata_args (list): Authorization metadata to pass to the grpc requests.
        Must contain 2 tuples with the function_id and authorization values.

    Returns:
      tuple[list, list]: A list of BlendShapeData and a list of blendshape names. The orders match.

    Notes:
    This function uploads the audio clip to the server, retrieves the hash,
    and then makes sequential FacePoseRequests to the server for each frame of the audio clip.

    """
    samplerate, data = scipy.io.wavfile.read(audio_clip)
    number_frames = int((len(data) / samplerate) * 30)
    # Create a secure gRPC channel and stub for the service.
    channel = auth.create_channel(uri=auth.GRPC_URI, use_ssl=True, metadata=metadata_args)
    stub = A2FAuthoringServiceStub(channel)
    (hash_gotten, bs_names) = upload_audio_clip_and_get_hash(stub, audio_clip)
    print("Perform sequential requests for the full audio clip...")
    req_list = [make_face_pose_request(hash_gotten, i * TIME_1_FRAME, bs_names) for i in range(number_frames)]
    bs_list = [get_avatar_face_pose(stub, req) for req in req_list]
    print("")
    return (bs_list, bs_names)


def create_parser():
    """
    Creates and returns an ArgumentParser object for a GRPC client tool with
    three subcommands: 'health_check', 'data_capture', and 'latency'.

    The parser allows users to perform various checks on the Authoring service,
    including checking service health, verifying data integrity, and
    evaluating performance and latency.

    Returns:
        argparse.ArgumentParser: An ArgumentParser object configured with the
        following subcommands:

        - health_check: Check the health status of the Authoring service.
            * api (str): The NGC API key to access the Catalog
            * function-id (str): The function id to invoke the API
            * version-id (str): The version id to invoke the API

        - data_capture: Capture data of generated by the Authoring service.
            * api (str): The NGC API key to access the Catalog
            * function-id (str): The function id to invoke the API
            * version-id (str): The version id to invoke the API
            * audio_clip (str): Path to the audio clip file for generating blendshape
                data
            * print_bs_names (bool): If true, the command will print out the blendshape names.

        - latency: Check the latency of the Authoring service.
            * api (str): The NGC API key to access the Catalog
            * function-id (str): The function id to invoke the API
            * version-id (str): The version id to invoke the API
            * audio_clip (str): Path to the audio clip file for testing.
            * number_requests (int): The number of requests to simulate.
            * concurrent (int): The number of concurrent requests to run.
            * print_bs_names (bool): If true, the command will print out the blendshape names.

    """
    parser = argparse.ArgumentParser(description="GRPC client tool connecting to NVCF deployment")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # Health check subparser
    health_check_parser = subparsers.add_parser("health_check", help="Check GRPC service health")
    health_check_parser.add_argument("--apikey", type=str, required=True, help="NGC API Key to invoke the API function")
    health_check_parser.add_argument(
        "--function-id", type=str, required=True, default="", help="Function ID to invoke the API function"
    )
    health_check_parser.add_argument(
        "--version-id", type=str, required=True, default="", help="Version ID to invoke the API function"
    )

    # Data capture check subparser
    data_capture_parser = subparsers.add_parser("data_capture", help="Check data integrity of GRPC service")
    data_capture_parser.add_argument("--apikey", type=str, required=True, help="NGC API Key to invoke the API function")
    data_capture_parser.add_argument(
        "--function-id", type=str, required=True, default="", help="Function ID to invoke the API function"
    )
    data_capture_parser.add_argument(
        "--version-id", type=str, required=True, default="", help="Version ID to invoke the API function"
    )
    data_capture_parser.add_argument("--audio-clip", required=True, help="Path to audio clip file")
    data_capture_parser.add_argument(
        "--print-bs-names",
        action="store_true",
        required=False,
        help="Optional. If enabled, the command will print out the names of the returned blendshapes.",
    )

    # Latency check subparser
    latency_parser = subparsers.add_parser("latency", help="Check latency of GRPC service")
    latency_parser.add_argument("--apikey", type=str, required=True, help="NGC API Key to invoke the API function")
    latency_parser.add_argument(
        "--function-id", type=str, required=True, default="", help="Function ID to invoke the API function"
    )
    latency_parser.add_argument(
        "--version-id", type=str, required=True, default="", help="Version ID to invoke the API function"
    )
    latency_parser.add_argument("--audio-clip", type=str, required=True, help="Path to audio clip file")
    latency_parser.add_argument(
        "--number-requests", type=positive_int, required=True, help="Number of requests to simulate"
    )
    latency_parser.add_argument("--concurrent", type=positive_int, required=True, help="Number of concurrent requests")
    latency_parser.add_argument(
        "--print-bs-names",
        action="store_true",
        required=False,
        help="Optional. If enabled, the command will print out the names of the returned blendshapes.",
    )

    return parser


def main():
    """Performs secure grpc calls to the Audio2Face-3D Authoring service hosted on NVCF."""
    parser = create_parser()
    args = parser.parse_args()

    # Authentication arguments are needed for all 3 modes of the script.
    metadata_args = [
        ("function-id", args.function_id),
        ("authorization", "Bearer " + args.apikey),
        ("version-id", args.version_id),
    ]

    if args.command == "health_check":
        # Checks the health of the service at the specified URL.
        # Prints "ONLINE" if the service is available, "OFFLINE" otherwise.
        is_healthy = check_health(auth.create_channel(uri=auth.GRPC_URI, use_ssl=True, metadata=metadata_args))
        print(f'NVCF Service is {"ONLINE" if is_healthy else "OFFLINE"}')

    elif args.command == "data_capture":
        (blendshape_data_list, blendshape_names) = get_blendshapes_for_each_frame(args.audio_clip, metadata_args)

        if args.print_bs_names:
            print(f"Blendshape names are: {blendshape_names}")
        with open(OUTPUT_BLENDSHAPE, "w", encoding="utf-8") as file_out:
            file_out.write("timecode," + ",".join(blendshape_names))
            file_out.write("\n")
            for i, blendshape_data in enumerate(blendshape_data_list):
                file_out.write(str(i * TIME_1_FRAME) + "," + ",".join(list(map(str, blendshape_data.blendshapes))))
                file_out.write("\n")
        print(f"Saved results to {OUTPUT_BLENDSHAPE}")
        with open(OUTPUT_EMOTIONS, "w", encoding="utf-8") as file_out:
            file_out.write("timecode," + ",".join(EMOTION_KEYS))
            file_out.write("\n")
            for i, blendshape_data in enumerate(blendshape_data_list):
                file_out.write(
                    str(i * TIME_1_FRAME)
                    + ","
                    + ",".join([str(blendshape_data.emotions[emotion]) for emotion in EMOTION_KEYS])
                )
                file_out.write("\n")
        print(f"Saved results to {OUTPUT_EMOTIONS}")

    elif args.command == "latency":
        # Computes latency data for the service and saves the results to a file.
        # Saves the latency data to a file named PERF_DATA_FILE and plots the data to an image
        # file named OUT_IMG.
        print(f"Computing data for {args.number_requests} requests with concurrency of " f"{args.concurrent}")

        channel_params = GrpcChannelParams(False, metadata_args=metadata_args)
        results = perform_parallel_data_exchange(args.audio_clip, args.number_requests, args.concurrent, channel_params)

        print()
        if args.print_bs_names:
            print(f"Blendshape names are: {results.blendshape_names}")
        with open(PERF_DATA_FILE, "w", encoding="utf-8") as file_perf_out:
            file_perf_out.write(str(results))
        results.plot_percentiles(OUT_IMG)
        print(f"Saved latency data in {PERF_DATA_FILE}")
        print(f"Plotted latency data in {OUT_IMG}")


if __name__ == "__main__":
    main()
